# -*- coding: utf-8 -*-
# @Time    : 2023/7/5 10:11
# @Author  : Pan
# @Software: PyCharm
# @Project : VisualFramework
# @FileName: Image2ImageDataset


import os
import paddle
import numpy as np
from paddle import io
from datasets import base
from transforms import Compose


class Image2ImageDataset(io.Dataset):
    def __init__(self, config):
        super(Image2ImageDataset, self).__init__()
        self.mode = config["mode"] if "mode" in config.keys() else "standard"
        self.gt_root = config["gt_root"] if "gt_root" in config.keys() else None
        self.dt_root = config["dt_root"] if "dt_root" in config.keys() else None
        self.mask_root = config["mask_root"] if "mask_root" in config.keys() else None
        self.trans = Compose(config["transforms"])

        self.data_list = self._make_list()

    def __getitem__(self, item):
        data = eval("self._"+self.mode)(item)
        return data

    def __len__(self):
        return len(self.data_list)

    def _standard(self, item):
        # print("||" + os.path.join(self.dt_root, self.data_list[item][0]) + "|" + os.path.join(self.gt_root, self.data_list[item][1]) + "||")
        data = {"img_1": os.path.join(self.dt_root, self.data_list[item]),
                "img_2": os.path.join(self.gt_root, self.data_list[item])}
        data = self.trans(data)
        return data

    def _predict(self, item):
        data = {"img": os.path.join(self.dt_root, self.data_list[item]),
                "path": self.data_list[item]}
        data = self.trans(data)
        return data

    def _make_list(self):
        data_list = []
        dt_list = os.listdir(self.dt_root)
        dt_list.sort()
        if self.mode == "standard":
            gt_list = os.listdir(self.gt_root)
            gt_list.sort()
            assert len(gt_list) == len(dt_list), "gt:%d与bt:%d数据数量不匹配" % (len(gt_list), len(dt_list))
            for i in range(len(dt_list)):
                assert gt_list[i] == gt_list[i], " %s %s:图像对不匹配" % (gt_list[i].split(".")[0], gt_list[i].split(".")[0])
                data_list.append(dt_list[i])
        else:
            for i in range(len(dt_list)):
                data_list.append(dt_list[i])
        return data_list
